# video/consumers.py
import json
from channels.generic.websocket import AsyncWebsocketConsumer
from faceapp.views import FaceViewSet
from faceapp.serializers import FaceSerializer
import base64
import cv2


class FaceConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        print(self.scope['url_route']['kwargs'])
        self.room_name = self.scope['url_route']['kwargs']['v_name']
        self.room_group_name = 'video_%s' % self.room_name
        # Join room group
        await self.channel_layer.group_add(
            self.room_group_name,
            self.channel_name
        )
        self.recognizer = FaceViewSet.recognizer
        await self.accept()

    async def disconnect(self, close_code):
        # Leave room group
        await self.channel_layer.group_discard(
            self.room_group_name,
            self.channel_name
        )

    # Receive message from WebSocket
    async def receive(self, text_data):
        print("### websocket 接受数据")
        # Send message to room group
        await self.channel_layer.group_send(
            self.room_group_name,
            {
                'type': 'video_message',
                'img': text_data,
            }
        )

    # Receive message from room group
    async def video_message(self, event):
        img = event['img']
        fea = self.recognizer.get_feature(img, align=False, base64img=True)
        if fea is None:
            await self.send(text_data=json.dumps({
                'det_img': "data:image/jpg;base64," +img,
                'reg_score': -1,
                'reg_img': ""
            }))
        else:
            res = self.recognizer.compare_1vn(fea)
            if res is None:
                await self.send(text_data=json.dumps({
                    'det_img': "data:image/jpg;base64," +img,
                    'reg_score': -1,
                    'reg_img': ""
                }))
            else:
                match_score = res['score']
                instance = res['instance']
                match_img = 'http://10.82.14.248:8000'+ instance.url.url

                # Send message to WebSocket
                await self.send(text_data=json.dumps({
                    'det_img': "data:image/jpg;base64," +img,
                    'reg_score': match_score,
                    'reg_img': match_img
                }))
